import json
import os
import glob
from elasticsearch import Elasticsearch
import hashlib
import time
import magic
import tika
from tika import parser
tika.initVM()
os.environ['JAVA_HOME'] = 'D:/jdk11'
os.environ['TIKA_STARTUP_SLEEP'] = '5'
os.environ['TIKA_STARTUP_MAX_RETRY'] = '3'
os.environ['TIKA_JAVA'] = 'D:/jdk11/bin/java.exe'  

# 提取文档标题的函数
def extract_document_title(parsed_result, filename):
    """
    从文档中提取标题
    策略：
    1. 尝试从元数据中获取标题
    2. 如果无元数据标题，从内容中提取第一个有意义的非空行
    3. 如果仍无标题，使用原始文件名
    """
    # 尝试从元数据获取标题
    metadata = parsed_result.get("metadata", {})
    title = metadata.get("title", "")
    
    # 如果元数据中没有标题，从内容中提取
    if not title and parsed_result.get("content"):
        content = parsed_result.get("content", "").strip()
        if content:
            lines = content.split('\n')
            # 查找第一个非空且有意义的行
            for line in lines:
                clean_line = line.strip()
                if clean_line and len(clean_line) >= 5 and len(clean_line) <= 200:  # 避免太短或太长的行
                    title = clean_line
                    break
    
    # 如果仍无标题，使用文件名（去除扩展名）
    if not title:
        title = os.path.splitext(os.path.basename(filename))[0]
    
    return title

# 连接到Elasticsearch
es = Elasticsearch(["http://localhost:9200"])

# 创建网页索引
def create_webpage_index():
    if es.indices.exists(index="nankai_webpages"):
        es.indices.delete(index="nankai_webpages")
    webpage_mapping = {
        "mappings": {
            "properties": {
                "url": {"type": "keyword"},
                "domain": {"type": "keyword"},
                "title": {"type": "text", "analyzer": "ik_max_word"},
                "text": {"type": "text", "analyzer": "ik_max_word"},
                "meta_keywords": {"type": "text", "analyzer": "ik_max_word"},
                "meta_description": {"type": "text", "analyzer": "ik_max_word"},
                "crawl_time": {"type": "date", "format": "yyyy-MM-dd HH:mm:ss"},
                "page_hash": {"type": "keyword"},
                "paragraphs": {"type": "text", "analyzer": "ik_max_word"},
                "links": {"type": "keyword"}
            }
        },
        "settings": {
            "index": {
                "number_of_shards": 5,
                "number_of_replicas": 1
            }
        }
    }
    es.indices.create(index="nankai_webpages", body=webpage_mapping)
    print("网页索引创建成功")

# 创建文档索引
def create_document_index():
    if es.indices.exists(index="nankai_documents"):
        es.indices.delete(index="nankai_documents")
    document_mapping = {
        "mappings": {
            "properties": {
                "url": {"type": "keyword"},
                "domain": {"type": "keyword"},
                "title": {"type": "text", "analyzer": "ik_max_word"},
                "file_type": {"type": "keyword"},
                "file_size": {"type": "long"},
                "crawl_time": {"type": "date", "format": "yyyy-MM-dd HH:mm:ss"},
                "content": {"type": "text", "analyzer": "ik_max_word"},
                "file_path": {"type": "keyword"},
                "content_type": {"type": "keyword"},
                "original_filename": {"type": "keyword"}
            }
        },
        "settings": {
            "index": {
                "number_of_shards": 5,
                "number_of_replicas": 1
            }
        }
    }
    es.indices.create(index="nankai_documents", body=document_mapping)
    print("文档索引创建成功")

# 处理网页数据
def index_webpage_data():
    # 使用清理后的数据目录
    files = glob.glob("crawler/cleaned_data/*.json")
    # 如果没有找到清理后的文件，回退到原始数据
    if not files:
        print("未找到清理后的数据，将使用原始数据")
        files = glob.glob("crawler/data/*.json")
    total_indexed = 0
    
    for file_path in files:
        print(f"处理文件: {file_path}")
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                try:
                    # 尝试加载整个文件作为一个JSON数组
                    all_data = json.load(f)
                    
                    # 如果是数组，则逐个处理其中的项目
                    if isinstance(all_data, list):
                        for data in all_data:
                            try:
                                # 只处理HTML页面，过滤文档
                                if 'file_type' not in data:
                                    # 生成ID避免重复
                                    doc_id = hashlib.md5(data['url'].encode()).hexdigest()
                                    
                                    # 格式化数据
                                    webpage_doc = {
                                        "url": data['url'],
                                        "domain": data['domain'],
                                        "title": data['title'],
                                        "text": data.get('text', ''),
                                        "meta_keywords": data.get('meta_keywords', ''),
                                        "meta_description": data.get('meta_description', ''),
                                        "crawl_time": data['crawl_time'],
                                        "page_hash": data.get('page_hash', ''),
                                        "paragraphs": data.get('paragraphs', []),
                                        "links": data.get('links', [])
                                    }
                                    
                                    # 索引文档
                                    es.index(index="nankai_webpages", id=doc_id, body=webpage_doc)
                                    total_indexed += 1
                                    
                                    if total_indexed % 1000 == 0:
                                        print(f"已索引 {total_indexed} 个网页")
                            except Exception as e:
                                print(f"处理数据项出错: {e}")
                                continue
                    else:
                        # 单个对象的情况
                        data = all_data
                        if 'file_type' not in data:
                            doc_id = hashlib.md5(data['url'].encode()).hexdigest()
                            # 格式化数据
                            webpage_doc = {
                                "url": data['url'],
                                "domain": data['domain'],
                                "title": data['title'],
                                "text": data.get('text', ''),
                                "meta_keywords": data.get('meta_keywords', ''),
                                "meta_description": data.get('meta_description', ''),
                                "crawl_time": data['crawl_time'],
                                "page_hash": data.get('page_hash', ''),
                                "paragraphs": data.get('paragraphs', []),
                                "links": data.get('links', [])
                            }
                            
                            # 索引文档
                            es.index(index="nankai_webpages", id=doc_id, body=webpage_doc)
                            total_indexed += 1
                
                except json.JSONDecodeError:
                    # 如果不是有效的JSON，尝试按行处理（可能是jsonlines格式）
                    f.seek(0)  # 回到文件开头
                    for line in f:
                        try:
                            data = json.loads(line.strip())
                            # 只处理HTML页面，过滤文档
                            if 'file_type' not in data:
                                # 生成ID避免重复
                                doc_id = hashlib.md5(data['url'].encode()).hexdigest()
                                
                                # 格式化数据
                                webpage_doc = {
                                    "url": data['url'],
                                    "domain": data['domain'],
                                    "title": data['title'],
                                    "text": data.get('text', ''),
                                    "meta_keywords": data.get('meta_keywords', ''),
                                    "meta_description": data.get('meta_description', ''),
                                    "crawl_time": data['crawl_time'],
                                    "page_hash": data.get('page_hash', ''),
                                    "paragraphs": data.get('paragraphs', []),
                                    "links": data.get('links', [])
                                }
                                
                                # 索引文档
                                es.index(index="nankai_webpages", id=doc_id, body=webpage_doc)
                                total_indexed += 1
                                
                                if total_indexed % 1000 == 0:
                                    print(f"已索引 {total_indexed} 个网页")
                        except Exception as e:
                            print(f"处理行数据出错: {e}")
        except Exception as e:
            print(f"处理文件 {file_path} 出错: {e}")
    
    print(f"网页索引完成，总共索引了 {total_indexed} 个网页")

# 提取文档内容并索引
def index_document_data():
    documents_dir = "crawler/data/documents"
    total_docs = 0
    
    # 从json文件获取文档元数据
    doc_metadata = {}
    # 使用清理后的数据目录
    files = glob.glob("crawler/cleaned_data/*.json")
    # 如果没有找到清理后的文件，回退到原始数据
    if not files:
        print("未找到清理后的文档元数据，将使用原始数据")
        files = glob.glob("crawler/data/*.json")
    
    for file_path in files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                try:
                    # 尝试加载整个文件作为一个JSON数组
                    all_data = json.load(f)
                    
                    # 如果是数组，则逐个处理其中的项目
                    if isinstance(all_data, list):
                        for data in all_data:
                            try:
                                if 'file_type' in data and 'file_path' in data:
                                    doc_metadata[data['file_path']] = data
                            except Exception as e:
                                print(f"处理元数据项出错: {e}")
                                continue
                    else:
                        # 单个对象的情况
                        data = all_data
                        if 'file_type' in data and 'file_path' in data:
                            doc_metadata[data['file_path']] = data
                
                except json.JSONDecodeError:
                    # 如果不是有效的JSON，尝试按行处理（可能是jsonlines格式）
                    f.seek(0)  # 回到文件开头
                    for line in f:
                        try:
                            data = json.loads(line.strip())
                            if 'file_type' in data and 'file_path' in data:
                                doc_metadata[data['file_path']] = data
                        except Exception as e:
                            pass
        except Exception as e:
            print(f"处理文件 {file_path} 元数据出错: {e}")
    
    # 处理文档文件
    for filename in os.listdir(documents_dir):
        file_path = os.path.join(documents_dir, filename)
        
        try:
            # 获取文件类型
            mime = magic.Magic(mime=True)
            file_type = mime.from_file(file_path)
            
            # 提取文本内容
            content = ""
            try:
                parsed = parser.from_file(file_path)
                content = parsed.get("content", "")
                # 清理内容
                if content:
                    content = content.strip()
                    
                # 提取文档标题
                doc_title = extract_document_title(parsed, filename)
                print(f"文件 {filename} 提取标题: {doc_title}")
                
            except Exception as e:
                print(f"无法解析文件 {filename}: {e}")
                doc_title = os.path.splitext(filename)[0]  # 使用文件名作为标题
            
            # 构建文档
            doc_id = hashlib.md5(filename.encode()).hexdigest()
            
            # 获取元数据（如果存在）
            metadata = doc_metadata.get(file_path, {})
            url = metadata.get('url', '')
            domain = metadata.get('domain', '')
            crawl_time = metadata.get('crawl_time', time.strftime('%Y-%m-%d %H:%M:%S'))
            
            doc = {
                "url": url,
                "domain": domain,
                "title": doc_title,  # 使用提取的标题
                "file_type": os.path.splitext(filename)[1][1:].lower(),
                "file_size": os.path.getsize(file_path),
                "crawl_time": crawl_time,
                "content": content,
                "file_path": file_path,
                "content_type": file_type,
                "original_filename": filename  # 保存原始文件名
            }
            
            # 索引文档
            es.index(index="nankai_documents", id=doc_id, body=doc)
            total_docs += 1
            
            if total_docs % 100 == 0:
                print(f"已索引 {total_docs} 个文档")
                
        except Exception as e:
            print(f"处理文档 {filename} 出错: {e}")
    
    print(f"文档索引完成，总共索引了 {total_docs} 个文档")

if __name__ == "__main__":
    # 创建索引
    create_webpage_index()
    create_document_index()
    
    # 索引数据
    print("开始索引网页数据...")
    index_webpage_data()
    
    print("开始索引文档数据...")
    index_document_data()
    
    print("所有数据索引完成！") 